# -*- coding: utf-8 -*-
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

from flyai.data_helper import DataHelper
from flyai.framework import FlyAI
from flyai.utils.log_helper import train_log

from path import *
from config import args, parser_info
from model import get_models, unfreeze_model
from confusion_matrix import confusion_matrix
from warm_scheduler import get_scheduler
from timer import Timer
from logger import logger
from losses import mse_loss as criterion_mse, FocalLoss
from ModelContainer import *

import torch
from torch.utils.data import DataLoader
from torch import nn
from torch.utils.tensorboard import SummaryWriter
from torch.optim.lr_scheduler import CosineAnnealingLR, LambdaLR

import numpy as np
import pandas as pd
import shutil


logger.logger.debug(parser_info)


'''
此项目为FlyAI2.0新版本框架，数据读取，评估方式与之前不同
2.0框架不再限制数据如何读取
样例代码仅供参考学习，可以自己修改实现逻辑。
模版项目下载支持 PyTorch、Tensorflow、Keras、MXNET、scikit-learn等机器学习框架
第一次使用请看项目中的：FlyAI2.0竞赛框架使用说明.html
使用FlyAI提供的预训练模型可查看：https://www.flyai.com/models
学习资料可查看文档中心：https://doc.flyai.com/
常见问题：https://doc.flyai.com/question.html
遇到问题不要着急，添加小姐姐微信，扫描项目里面的：FlyAI小助手二维码-小姐姐在线解答您的问题.png
'''
if not os.path.exists(MODEL_PATH):
    os.makedirs(MODEL_PATH)


class Main(FlyAI):
    def __init__(self):
        pass

    def download_data(self):
        data_helper = DataHelper()
        data_helper.download_from_ids("BaldClassification")

    def train(self):
        import ModelEnsemble


if __name__ == '__main__':
    if os.path.exists(LOG_PATH):
        shutil.rmtree(LOG_PATH)
    main = Main()
    if args.train_flyai:
        main.download_data()
    main.train()